import pandas as pd


def cumulative_to_monthly(df_cum, value_columns=None):
    """
    将时间序列中的月度累计值转换为月度当月值
    
    参数:
        df_cum: DataFrame，索引为 datetime 格式，包含需要转换的累计值列
        value_columns: 列表，指定需要转换的列名。若为 None，则转换所有数值列
    
    返回:
        DataFrame: 转换后的月度当月值数据
    """
    df_cum = df_cum.sort_index()
    # 复制数据避免修改原数据
    df = df_cum.copy(deep=True)
    
    # 确保索引是 datetime 格式
    if not pd.api.types.is_datetime64_any_dtype(df.index):
        df.index = pd.to_datetime(df.index)
    
    # 确定需要转换的列
    if value_columns is None:
        value_columns = df.select_dtypes(include=['number']).columns.tolist()
    
    # 对指定列进行累计值转当月值计算
    for col in value_columns:
        # 计算当月值 = 当前累计值 - 上月累计值
        df[col] = df_cum[col] - df_cum[col].shift(1)
        
        # 处理每年1月的数据（1月当月值等于其累计值，因为累计从每年1月重新开始）
        # 提取年份和月份信息
        df['year'] = df.index.year
        df['month'] = df.index.month
        
        # 每年1月的当月值等于其累计值（因为没有上一个月的累计值）
        is_first_month = df['month'] == 1
        df.loc[is_first_month, col] = df_cum.loc[is_first_month, col]
        
    # 删除临时列
    df.drop(columns=['year', 'month'], inplace=True)
    
    return df


def monthly_to_quarterly(df_monthly, value_columns=None, agg_method='mean'):
    """
    将月度频率数据转换为季度频率数据
    
    参数:
        df_monthly: DataFrame，索引为 datetime 格式，包含月度数据
        value_columns: 列表，指定需要转换的列名。若为 None，则转换所有数值列
        agg_method: 聚合方式，可选值为 'mean'（平均）、'sum'（求和）、'first'（第一个值）、'last'（最后一个值）
    
    返回:
        DataFrame: 转换后的季度频率数据，索引保持datetime格式
    """
    # 复制数据避免修改原数据
    df = df_monthly.copy(deep=True)
    
    # 确保索引是 datetime 格式
    if not pd.api.types.is_datetime64_any_dtype(df.index):
        df.index = pd.to_datetime(df.index)
    
    # 确定需要转换的列
    if value_columns is None:
        value_columns = df.select_dtypes(include=['number']).columns.tolist()
    
    # 检查聚合方法是否合法
    valid_methods = ['mean', 'sum', 'first', 'last']
    if agg_method not in valid_methods:
        raise ValueError(f"聚合方法 {agg_method} 不支持，请使用 {valid_methods} 中的一种")
    
    # 按季度重采样并聚合，保持索引为datetime格式
    # 使用 Q-DEC 表示以12月为年末的季度（Q1:1-3月, Q2:4-6月, ..., Q4:10-12月）
    quarterly_df = df[value_columns].resample('Q-DEC').agg(agg_method)
    
    # 移除了将索引转换为period格式的操作，保持原生datetime索引
    # 索引将默认使用季度末的日期（如3月31日、6月30日等）作为datetime标识
    
    return quarterly_df


def calculate_yoy(df, value_columns=None):
    """
    计算时间序列数据的同比增长率，兼容月度和季度数据
    
    参数:
        df: DataFrame，索引为 datetime 格式，包含需要计算同比的数据
        value_columns: 列表，指定需要计算的列名。若为 None，则计算所有数值列
    
    返回:
        DataFrame: 包含同比增长率的DataFrame，索引与输入保持一致，
                  列名为原列名后加 '_yoy' 后缀
    """
    # 复制数据避免修改原数据
    df = df.copy(deep=True)
    
    # 确保索引是 datetime 格式
    if not pd.api.types.is_datetime64_any_dtype(df.index):
        df.index = pd.to_datetime(df.index)
    
    # 确定需要计算的列
    if value_columns is None:
        value_columns = df.select_dtypes(include=['number']).columns.tolist()
    
    # 创建结果DataFrame
    result = pd.DataFrame(index=df.index)
    
    # 判断数据频率（月度或季度）
    # 通过检查时间间隔模式判断
    if len(df) >= 2:
        freq = pd.infer_freq(df.index)
        if freq is None:
            raise ValueError("无法推断数据频率，请确保输入数据为规则的月度或季度时间序列")
        
        # 确定偏移量（月度为12个月，季度为4个季度）
        if freq.startswith(('M', 'MS')):  # 月度频率
            shift_periods = 12
        elif freq.startswith(('Q', 'QS')):  # 季度频率
            shift_periods = 4
        else:
            raise ValueError("仅支持月度和季度数据的同比计算")
    else:
        raise ValueError("数据量不足，无法判断频率")
    
    # 计算同比增长率：(当前值/上年同期值 - 1) * 100%
    for col in value_columns:
        # 获取上年同期值
        prev_year_values = df[col].shift(shift_periods)
        # 计算同比（避免除零错误）
        result[f'{col}_yoy'] = (df[col] / prev_year_values - 1) * 100
    
    return result